import os

import numpy as np
import pandas as pd

base_path = os.path.dirname(__file__)


class GetData:
    def getData(self):
        self.source_vars_name = ['alpha', 'CL', 'CD', 'CDp', 'Re', 'Cl/Cd']

        self.source_x_name = ['alpha']
        self.source_x_mean = ['α']

        self.source_y_name = ['CL', 'CD']
        self.source_y_mean = ['Cl', 'Cd']

        self.target_vars_name = ['alpha', 'Cl', 'Cd', 'Re', 'Cl/Cd']

        self.target_x_name = ['alpha']
        self.target_x_mean = ['α']

        self.target_y_name = ['Cl', 'Cd']
        self.target_y_mean = ['Cl', 'Cd']

        self.source_data = pd.read_excel(f"{base_path}/NACA 2414/Re=60400/XFLR_Re60400.xlsx", engine='openpyxl')
        self.target_data = pd.read_excel(f"{base_path}/NACA 2414/Re=60400/UIUC_CS_Re60400.xlsx", engine='openpyxl')

        self.source_x = self.source_data[self.source_x_name].to_numpy()
        self.source_y = self.source_data[self.source_y_name].to_numpy()

        self.target_x = self.target_data[self.target_x_name].to_numpy()
        self.target_y = self.target_data[self.target_y_name].to_numpy()

        # 训练集占比===================================================
        # source默认1/2
        self.source_x_train = self.source_x[0: self.source_x.shape[0]: 2]
        self.source_x_test = self.source_x[1: self.source_x.shape[0]: 2]
        self.source_y_train = self.source_y[0: self.source_y.shape[0]: 2]
        self.source_y_test = self.source_y[1: self.source_y.shape[0]: 2]

        # target
        # 1/2=============================
        # self.target_x_train = self.target_x[0: self.target_x.shape[0]: 2]
        # self.target_x_test = self.target_x[1: self.target_x.shape[0]: 2]
        #
        # self.target_y_train = self.target_y[0: self.target_y.shape[0]: 2]
        # self.target_y_test = self.target_y[1: self.target_y.shape[0]: 2]

        # 1/3================================
        # self.target_x_train = self.target_x[0:self.target_x.shape[0]:3]
        # x_test1 = self.target_x[1:self.target_x.shape[0]:3]
        # x_test2 = self.target_x[2:self.target_x.shape[0]:3]
        # self.target_x_test = np.concatenate((x_test1, x_test2))
        #
        # self.target_y_train = self.target_y[0:self.target_y.shape[0]:3]
        # y_test1 = self.target_y[1:self.target_y.shape[0]:3]
        # y_test2 = self.target_y[2:self.target_y.shape[0]:3]
        # self.target_y_test = np.concatenate((y_test1, y_test2))

        # 1/4================================
        # self.target_x_train = self.target_x[0:self.target_x.shape[0]:4]
        # x_test1 = self.target_x[1:self.target_x.shape[0]:4]
        # x_test2 = self.target_x[2:self.target_x.shape[0]:4]
        # x_test3 = self.target_x[3:self.target_x.shape[0]:4]
        # self.target_x_test = np.concatenate((x_test1, x_test2, x_test3))
        #
        # self.target_y_train = self.target_y[0:self.target_y.shape[0]:4]
        # y_test1 = self.target_y[1:self.target_y.shape[0]:4]
        # y_test2 = self.target_y[2:self.target_y.shape[0]:4]
        # y_test3 = self.target_y[3:self.target_y.shape[0]:4]
        # self.target_y_test = np.concatenate((y_test1, y_test2, y_test3))

        # 1/5================================
        self.target_x_train = self.target_x[0:self.target_x.shape[0]:5]
        x_test1 = self.target_x[1:self.target_x.shape[0]:5]
        x_test2 = self.target_x[2:self.target_x.shape[0]:5]
        x_test3 = self.target_x[3:self.target_x.shape[0]:5]
        x_test4 = self.target_x[4:self.target_x.shape[0]:5]
        self.target_x_test = np.concatenate((x_test1, x_test2, x_test3, x_test4))

        self.target_y_train = self.target_y[0:self.target_y.shape[0]:5]
        y_test1 = self.target_y[1:self.target_y.shape[0]:5]
        y_test2 = self.target_y[2:self.target_y.shape[0]:5]
        y_test3 = self.target_y[3:self.target_y.shape[0]:5]
        y_test4 = self.target_y[4:self.target_y.shape[0]:5]
        self.target_y_test = np.concatenate((y_test1, y_test2, y_test3, y_test4))

        # # 排序 ====================================================
        all_data = np.hstack((self.target_x_test, self.target_y_test))  # 合并数据集
        sorted_indices = np.argsort(all_data[:, 0])  # 排序规则
        all_data = all_data[sorted_indices]  # 排序
        self.target_x_test = all_data[:, :self.target_x_test.shape[1]]
        self.target_y_test = all_data[:, self.target_x_test.shape[1]:]


if __name__ == '__main__':
    data = GetData()
    data.getData()
